from segment_anything import build_sam, SamPredictor 
from segment_anything import sam_model_registry, SamPredictor
import os,pdb,torch,cv2
import numpy as np
from PIL import Image
'''
efficientsam_s.pth   mobile_sam.pth        sam_vit_h_4b8939.pth
efficientsam_ti.pth  sam_vit_b_01ec64.pth  sam_vit_l_0b3195.pth
'''
SAM_CKP = '/home/shengjie/ckp/segment-anything/checkpoints/sam_vit_h_4b8939.pth'
examples_dir = '/data/shengjie/style_zhenzhi/'
save_dir = '/data/shengjie/synthesis_zhenzhi/'

# 初始化模型
# sam_checkpoint = "sam_vit_h_4b8939.pth"  # 模型路径
model_type = "vit_h"  # 模型类型(vit_h/vit_l/vit_b)
device = "cuda" if torch.cuda.is_available() else "cpu"  # 自动选择设备

# 加载模型
sam = sam_model_registry[model_type](checkpoint=SAM_CKP)
sam.to(device=device)
predictor = SamPredictor(sam)

# pdb.set_trace()

# imagefiles = os.listdir(examples_dir)

# test_img = os.path.join(examples_dir,imagefiles[0])
# prompt='cloth'

def get_mask_by_sam(image_path:str=None,image_pil=None,image_np=None):
    # 加载图像
    if image_path is not None:
        image = cv2.imread(image_path)  # 替换为你的图片路径
    elif image_pil is not None:
        image = np.array(image_pil)
    elif image_np is not None:
        image = image_np
    else:
        raise Exception('传入的图像数据不合理，无法获得mask')
    h,w,_ = image.shape
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为RGB格式
    predictor.set_image(image)  # 预处理图像

    # 定义提示点（前景点+背景点）
    input_point = np.array([[h//2, w//2], [0, 0]])  # 替换为你的坐标
    input_label = np.array([1, 0])  # 1=前景, 0=背景


    # 执行预测（可替换为框提示）
    masks, scores, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=True  # 返回多个候选结果
    )


    # pdb.set_trace()

    index = np.argmax(scores)
    mask = masks[index].astype(np.uint8)*255
    # Image.fromarray(mask).save(f'mask-{index}.png')
    mask_pil = Image.fromarray(mask)
    return mask,mask_pil
